import requests


def siliconflow_reranker(query:str,documents:list[str],top_n:int):
    url = "https://api.siliconflow.cn/v1/rerank"

    payload = {
        "model": "BAAI/bge-reranker-v2-m3",
        "query": query,
        "documents": documents,
        "top_n": top_n,
        # "return_documents": "false" , # 是否返回文本
        "max_chunks_per_doc": 1024,
        "overlap_tokens": 79
    }
    headers = {
        "Authorization": "Bearer sk-nihhjskdaommndqhohdrneftvrfukgavarvteuearcrdjgqw",
        "Content-Type": "application/json"
    }
    response = requests.request("POST", url, json=payload, headers=headers)
    # {"id":"019363fd560677a49bfd180cdc7f6abb","results":[{"index":1,"relevance_score":0.9953725},{"index":2,"relevance_score":0.002157342},{"index":0,"relevance_score":0.00046371284},{"index":3,"relevance_score":0.000017502925}],"meta":{"billed_units":{"input_tokens":28,"output_tokens":0,"search_units":0,"classifications":0},"tokens":{"input_tokens":28,"output_tokens":0}}}
    index_list = response.json()["results"]
    index_list = [item["index"] for item in index_list]
    return index_list

if __name__ == '__main__':
    query = "Hello, my name is John."
    documents = ["Hello, my name is John.", "I am a student.","I am a student.","I am a student."]
    index_list = siliconflow_reranker(query,documents,30)
    print(index_list)